Fix FSDP2 defaulting to version 1 in TrainingArguments; add dynamic plugin param passthrough by amanzoni1 · Pull Request #42521 · huggingface/transformers

amanzoni1 · 2025-12-01T16:14:18Z

What does this PR do?

Fixes the bug where TrainingArguments(fsdp=True, fsdp_config={"fsdp_version": 2, ...}) defaults to FSDP version 1, ignoring the version (unless the Accelerator was initialized manually), and most params were lost.
Adds dynamic passthrough of all FSDP plugin params (FullyShardedDataParallelPlugin) from fsdp_config to the plugin args (future-proof, no hardcoding).

Changes

Extract fsdp_version pre-stripping to preserve it.
Skip stripping for "fsdp_version" to avoid mangling.
Dynamic copy from FullyShardedDataParallelPlugin dataclass for all fields.
Add test_fsdp_version_2_config for regression.

Repro

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="/tmp/test",
    fsdp=True,
    fsdp_config={
        "fsdp_version": 2,
        "reshard_after_forward": True,
    },
)

# Call the internal method to test argument parsing
plugin_args = args._process_fsdp_args()

# Expected output after this PR:
print(f"fsdp_version: {plugin_args['fsdp_version']}")                      # Before: 1, After: 2
print(f"reshard_after_forward: {plugin_args['reshard_after_forward']}")    # Before: None, After: True

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
[ X ] Did you read the contributor guideline,
Pull Request section?
[ X ] Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case.
https://discuss.huggingface.co/t/how-to-start-fsdp2-when-using-trainer/151885
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
[ X ] Did you write any new necessary tests?

Who can review?

@SunMarc @3outeille

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

SunMarc

Thanks ! Left a few comments

SunMarc · 2025-12-02T14:22:03Z

-            fsdp_plugin_args["fsdp_version"] = self.fsdp_config.get("fsdp_version", 1)
-            prefetch_policy = self.fsdp_config.get("backward_prefetch", "NO_PREFETCH")
-            fsdp_plugin_args["backward_prefetch"] = prefetch_policy.upper()
-            fsdp_plugin_args["forward_prefetch"] = str(self.fsdp_config.get("forward_prefetch", "false")).lower()
-
-            sync_module_states = str(self.fsdp_config.get("sync_module_states", "true")).lower()
-            cpu_ram_efficient_loading = str(self.fsdp_config.get("cpu_ram_efficient_loading", "false")).lower()
-            if sync_module_states == "false" and cpu_ram_efficient_loading == "true":
-                # In this case, all the processes except the main process would have random weights leading
-                # to unexpected behaviour during training, thus throwing error here to prevent it.
-                raise ValueError('`sync_module_states` must be `"True"` if `cpu_ram_efficient_loading` is `"True"`')
-
-            # we need to set the env here as otherwise we get a warning in accelerate + we need to set it for transformers
-            fsdp_plugin_args["cpu_ram_efficient_loading"] = cpu_ram_efficient_loading
-            os.environ["FSDP_CPU_RAM_EFFICIENT_LOADING"] = cpu_ram_efficient_loading
-
-            fsdp_plugin_args["sync_module_states"] = sync_module_states
-            fsdp_plugin_args["use_orig_params"] = str(self.fsdp_config.get("use_orig_params", "true")).lower()


I think it will be better to keep this for now as we can easily change default config here instead of doing it in accelerate. Also, some defaults are not the same here vs accelerate. Nevertheless, for keys that are not exposed here but we have them in FullyShardedDataParallelPlugin, it's fine to set them.

SunMarc · 2025-12-02T14:22:44Z

+                if key in self.fsdp_config and key not in fsdp_plugin_args:
+                    fsdp_plugin_args[key] = self.fsdp_config[key]


we can keep this but put it at the end.

SunMarc · 2025-12-02T14:22:50Z

+                    # Set deprecated sharding_strategy from CLI (plugin maps to reshard_after_forward)
+                    # Skip if config has explicit reshard_after_forward (prioritize config)
+                    if "reshard_after_forward" not in self.fsdp_config:
+                        fsdp_plugin_args["sharding_strategy"] = fsdp_option


SunMarc · 2025-12-02T14:23:00Z

        if self.fsdp_config is not None and isinstance(self.fsdp_config, dict):
            for k in list(self.fsdp_config.keys()):
-                if k.startswith("fsdp_"):
+                if k.startswith("fsdp_") and k != "fsdp_version":


indeed thanks !

SunMarc · 2025-12-02T14:23:11Z

+    def test_fsdp_version_2_config(self):
+        output_dir = self.get_auto_remove_tmp_dir()
+        kwargs = {
+            "output_dir": output_dir,
+            "train_len": 128,
+            "save_steps": 5,
+            "learning_rate": 0.1,
+            "fsdp": True,
+            "fsdp_config": {
+                "fsdp_version": 2,
+                "reshard_after_forward": True,
+            },
+        }
+        with mockenv_context(**self.dist_env_1_gpu):
+            trainer = get_regression_trainer(**kwargs)
+            plugin_args = trainer.args._process_fsdp_args()
+            self.assertEqual(plugin_args["fsdp_version"], 2)
+            self.assertTrue(plugin_args["reshard_after_forward"])
+


amanzoni1 · 2025-12-02T17:41:37Z

Thanks for the feedback @SunMarc!!
implemented the hardcoded defaults + dynamic copy at end, and needed to add a conditional skip for forward_prefetch in v2 (Accelerate raises ValueError otherwise: "ValueError: forward_prefetch is not yet implemented in FSDP2, set to None or use fsdp_version=1").
Let me know if that can works

SunMarc · 2025-12-02T18:06:13Z

I've merged this PR that seems to fix a bit more things ! Feel free to rebase and add your updates, especially the test is nice !

#42566

amanzoni1 · 2025-12-02T20:31:53Z

@SunMarc Thanks for the feedback! Rebased onto the latest main, and ready for review

SunMarc

Thanks, just a few nits !

SunMarc · 2025-12-03T13:04:38Z

+            # HF-to-plugin map
+            if (
+                "transformer_layer_cls_to_wrap" in self.fsdp_config
+                and "transformer_cls_names_to_wrap" not in fsdp_plugin_args
+            ):
+                fsdp_plugin_args["transformer_cls_names_to_wrap"] = ",".join(
+                    self.fsdp_config["transformer_layer_cls_to_wrap"]
+                )


not need anymore

SunMarc · 2025-12-03T13:06:05Z

+            "fsdp_config": {
+                "fsdp_version": 2,
+                "reshard_after_forward": True,
+            },


can you add more things in the config so that it is better tested ?

SunMarc

Thanks !

HuggingFaceDocBuilderDev · 2025-12-03T16:01:17Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

amanzoni1 · 2025-12-08T12:23:00Z

Hey @SunMarc, not sure if there's anything more I should do here. The failed CI doesn't seem code related. Thanks!

SunMarc · 2025-12-08T13:25:20Z

cc @ArthurZucker can you merge ?

github-actions · 2025-12-08T13:32:40Z

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=42521&sha=661f54

github-actions · 2025-12-08T13:49:58Z

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=42521&sha=dd3f32

huggingface#45640) * FSDP: default to v2, simplify API around fsdp + fsdp_config - TrainingArguments.fsdp is now a boolean on-switch. String / list values are still accepted for backward compatibility but are translated into fsdp_config entries (and emit a deprecation warning). - fsdp_config now defaults to FSDP2 (version=2). FSDP1 usage still works but logs a deprecation warning (to be removed in v5.20). - New fsdp_config keys: auto_wrap_policy (default TRANSFORMER_BASED_WRAP), cpu_offload, state_dict_type (default FULL_STATE_DICT so trainer.save_model() produces an HF-compatible checkpoint out of the box). - FSDP1-only handling (forward_prefetch, backward_prefetch, use_orig_params, sync_module_states, and the string form of reshard_after_forward) is now isolated in a single branch for easy removal once v1 support is dropped. Legacy `fsdp` string/list parsing lives in a single _apply_legacy_fsdp_to_config helper for the same reason. - Trainer reads args.fsdp / args.fsdp_config guarded so that running without FSDP (or via `accelerate launch` with no transformers-side config) no longer crashes on None. - Docstring rewritten around fsdp_config; FSDP2-only keys surfaced, FSDP1-only keys (sync_module_states, use_orig_params, limit_all_gathers) dropped from the public docstring. Verified on the repro from PR huggingface#42521: TrainingArguments(fsdp=True, fsdp_config={"fsdp_version": 2, "reshard_after_forward": True}) now produces fsdp_version=2 and reshard_after_forward=True. * Apply repo consistency fixes * Cleanup: drop unused bool() casts and stale comment - `args.fsdp` is already a bool after `_process_fsdp_args`, so wrapping it with `bool(...)` in trainer.py was redundant. - Remove the explanatory comment above `state_dict_type` default; the docstring already covers it. * Forward activation_checkpointing / limit_all_gathers via plugin kwargs Drop the post-accelerator `setattr` loop that re-pushed these two keys onto the already-constructed FSDP plugin. They are now forwarded during plugin construction via `fsdp_plugin_args`: - `activation_checkpointing` is shared between FSDP1 and FSDP2. - `limit_all_gathers` is FSDP1-only (obsolete in FSDP2), so it lives in the v1 branch. The `activation_checkpointing` + `gradient_checkpointing` conflict check stays in trainer.py (still needs the post-plugin state to compare). * Apply suggestions from code review Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * fix * Apply repo consistency fixes --------- Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

ArthurZucker · 2026-10-20T12:13:15Z

Thanks a lot for the fix and your patience, @amanzoni1, and sorry this didn't get merged in time! Since then, #45640 rewrote how fsdp_config is turned into the Accelerate plugin and made FSDP2 the default (fsdp_config["version"] defaults to 2), so the version being ignored is fixed on main. tests/fsdp/test_fsdp.py was also replaced by tests/trainer/distributed/ in #44338. If some plugin options you needed are still not passed through, a small PR against main with a test under tests/trainer/distributed/ would be very welcome. Closing this one, thanks again!

amanzoni1 force-pushed the fix-fsdp2-default-version branch 2 times, most recently from f19d66a to bcd3599 Compare December 2, 2025 11:16

SunMarc reviewed Dec 2, 2025

View reviewed changes

amanzoni1 force-pushed the fix-fsdp2-default-version branch from 9a8af59 to b6a9f0b Compare December 2, 2025 20:10

SunMarc reviewed Dec 3, 2025

View reviewed changes

amanzoni1 force-pushed the fix-fsdp2-default-version branch from 61c0e74 to 2252af4 Compare December 3, 2025 15:33

SunMarc approved these changes Dec 3, 2025

View reviewed changes

SunMarc enabled auto-merge (squash) December 3, 2025 15:50

amanzoni1 added 4 commits December 3, 2025 20:10

Fix FSDP v2 defaulting to version 1 in TrainingArguments

139cb6b

Fix FSDP2 params and add test

f14bd4f

Fix FSDP2 params and add test

1140f69

Better FSDP2 test

ff47d0d

auto-merge was automatically disabled December 3, 2025 16:10
Head branch was pushed to by a user without write access

amanzoni1 force-pushed the fix-fsdp2-default-version branch from 3761cdd to ff47d0d Compare December 3, 2025 16:10

Merge branch 'main' into fix-fsdp2-default-version

7a0a30e

SunMarc enabled auto-merge (squash) December 3, 2025 16:53

Merge branch 'main' into fix-fsdp2-default-version

661f54d

Merge branch 'main' into fix-fsdp2-default-version

dd3f327

This was referenced Apr 29, 2026

Cumulative feature and defect updates from recent Transformers PRs evalstate/transformers#42

Open

Cumulative defect fixes from recent Transformers PRs evalstate/transformers#43

Open

ArthurZucker closed this Oct 20, 2026

auto-merge was automatically disabled October 20, 2026 12:13
Pull request was closed

amanzoni1 mentioned this pull request Nov 11, 2026

Forward mixed_precision_policy, ignored_modules and activation_checkpointing_offload from fsdp_config to the FSDP plugin #50628

Open

		if key in self.fsdp_config and key not in fsdp_plugin_args:
		fsdp_plugin_args[key] = self.fsdp_config[key]

Conversation

amanzoni1 commented Dec 1, 2025

What does this PR do?

Changes

Repro

Before submitting

Who can review?

Uh oh!

SunMarc left a comment

Choose a reason for hiding this comment

Uh oh!

SunMarc Dec 2, 2025

Choose a reason for hiding this comment

Uh oh!

SunMarc Dec 2, 2025

Choose a reason for hiding this comment

Uh oh!

SunMarc Dec 2, 2025

Choose a reason for hiding this comment

Uh oh!

SunMarc Dec 2, 2025

Choose a reason for hiding this comment

Uh oh!

SunMarc Dec 2, 2025

Choose a reason for hiding this comment

Uh oh!

amanzoni1 commented Dec 2, 2025

Uh oh!

SunMarc commented Dec 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

amanzoni1 commented Dec 2, 2025

Uh oh!

SunMarc left a comment

Choose a reason for hiding this comment

Uh oh!

SunMarc Dec 3, 2025

Choose a reason for hiding this comment

Uh oh!

SunMarc Dec 3, 2025

Choose a reason for hiding this comment

Uh oh!

SunMarc left a comment

Choose a reason for hiding this comment

Uh oh!

HuggingFaceDocBuilderDev commented Dec 3, 2025

Uh oh!

amanzoni1 commented Dec 8, 2025

Uh oh!

SunMarc commented Dec 8, 2025

Uh oh!

github-actions Bot commented Dec 8, 2025

Uh oh!

github-actions Bot commented Dec 8, 2025

Uh oh!

ArthurZucker commented Oct 20, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

SunMarc commented Dec 2, 2025 •

edited

Loading